import os
import sys
import glob
import numpy as np

import pydicom
from pydicom.tag import Tag

PHOTOMETRIC_INTERPRETATION_TAG = (0x0028,0x0004)
BITS_STORED_TAG = (0x0028, 0x0101)
RESCALE_SLOPE_TAG = (0x0028,0x1053)
RESCALE_INTERCEPT_TAG = (0x0028,0x1052)
BODY_PART_TAG = (0x0018, 0x0015)

DR_DICTS = ('DR', 'DX', 'CR', 'RF', 'XA', 'OT')
BODY_PART_NAME = ('LSPINE', 'LSPINE', 'SKULL', 'ANKLE', 'HAND', 'KNEE',
                  'ABDOMEN', 'CSPINE', 'SHOULDER', 'FOOT', 'CHEST', 'HIP', 'THORAX', '')

sys_is_little_endian = (sys.byteorder == 'little')

NumpySupportedTransferSyntaxes = [
    pydicom.uid.ExplicitVRLittleEndian,
    pydicom.uid.ImplicitVRLittleEndian,
    pydicom.uid.DeflatedExplicitVRLittleEndian,
    pydicom.uid.ExplicitVRBigEndian,
]

def should_change_PhotometricInterpretation_to_RGB(dicom_dataset):
    return False

def dicom_convert_jpg(dicom_path, jpg_path):
    if not os.path.exists(dicom_path):
        return -1

    ds = pydicom.dcmread(dicom_path)

    # Convert to float to avoid overflow or underflow losses.
    image_2d = ds.pixel_array.astype(float)

    # Rescaling grey scale between 0-255
    image_2d_scaled = (np.maximum(image_2d, 0) / image_2d.max()) * 255.0

    # Convert to uint
    image_2d_scaled = np.uint8(image_2d_scaled)
    #io.imsave(jpg_path, image_2d_scaled)

    return 0


def create_hu_lookup(bits):
    lookup = [i for i in range(1<<bits)]
    return lookup


def read_CR_dcm(fileanme):
    dcm_info = {}

    if not os.path.exists(fileanme):
        return dcm_info

    #使用context manager打开比较安全
    #with pydicom.dcmread("rtplan.dcm") as ds:
    #  ds.PatientName
    #todo
    dcm = pydicom.dcmread(fileanme)
    dcm.decompress()

    dcm_info['Modality'] = dcm.Modality

    if Tag(BODY_PART_TAG) in dcm:
        body_name = dcm.BodyPartExamined
        if body_name not in BODY_PART_NAME:
            encode_name = body_name.encode(encoding="iso8859")
            body_name = encode_name.decode('gbk')
        dcm_info['body_part'] = body_name # str(dcm[BODY_PART_TAG[0], BODY_PART_TAG[1]]._value).upper()
    else:
        dcm_info['body_part'] = ''

    if Tag(PHOTOMETRIC_INTERPRETATION_TAG) in dcm:
        dcm_info['photo_metric_interpretation'] = dcm[PHOTOMETRIC_INTERPRETATION_TAG[0], PHOTOMETRIC_INTERPRETATION_TAG[1]]._value
    else:
        dcm_info['photo_metric_interpretation'] = ''

    if Tag(BITS_STORED_TAG) in dcm:
        dcm_info['bits'] = int(dcm[BITS_STORED_TAG[0], BITS_STORED_TAG[1]]._value)
    else:
        dcm_info['bits'] = 8

    if Tag(RESCALE_SLOPE_TAG) in dcm:
        dcm_info['rescaleslope'] = int(dcm[RESCALE_SLOPE_TAG[0], RESCALE_SLOPE_TAG[1]]._value)
    else:
        dcm_info['rescaleslope'] = 0

    if Tag(RESCALE_INTERCEPT_TAG) in dcm:
        dcm_info['rescaleintercept'] = int(dcm[RESCALE_INTERCEPT_TAG[0], RESCALE_INTERCEPT_TAG[1]]._value)
    else:
        dcm_info['rescaleintercept'] = 1

    img_data = dcm.pixel_array
    dcm_info['rows'] = int(dcm.Rows)
    dcm_info['cols'] = int(dcm.Columns)

    if hasattr(dcm, 'WindowCenter'):
        dcm_info['wc'] = int(dcm.WindowCenter)
    else:
        dcm_info['wc'] = 0
    if hasattr(dcm, 'WindowWidth'):
        dcm_info['ww'] = int(dcm.WindowWidth)
    else:
        dcm_info['ww'] = 0

    img_type = dcm_info['Modality'].upper()
    if hasattr(dcm, 'pixel_array'):
        img_data = dcm.pixel_array
        # dcm_info['dtype'] = img_data.dtype.name
        # img_data = img_data.astype(float)
        #
        # if img_type in DR_DICTS:
        #     if dcm_info['wc'] != 0 and dcm_info['ww'] != 0:
        #         img_data = setDicomWinWidthWinCenter(img_data, dcm_info)
        #     else:
        #         img_data = convert_defalut_ww_wc(img_data)
        # else:
        #     img_data = convert_defalut_ww_wc(img_data)
        #
        # if dcm_info['photo_metric_interpretation'].upper().find('MONOCHROME1') >=0:
        #     img_data = (255 - img_data)
    else:
        img_data = None
    dcm_info['image'] = img_data
    return dcm_info


def convert_dcm_grayscale8(dcm_info):
    img_data = dcm_info['image'].astype(float)
    img_type = dcm_info['Modality'].upper()

    if img_type in DR_DICTS:
        if dcm_info['wc'] != 0 and dcm_info['ww'] != 0:
            img_data = setDicomWinWidthWinCenter(img_data, dcm_info['ww'], dcm_info['wc'])
        else:
            img_data = convert_defalut_ww_wc(img_data)
    else:
        img_data = convert_defalut_ww_wc(img_data)

    if dcm_info['photo_metric_interpretation'].upper().find('MONOCHROME1') >= 0:
        img_data = (255 - img_data)

    return img_data

def ajuste_dcm_grayscale8(dcm_info, ww, wc):
    img_data = dcm_info['image'].astype(float)
    img_type = dcm_info['Modality'].upper()

    if img_type in DR_DICTS:
        if ww != 0 and wc != 0:
            img_data = setDicomWinWidthWinCenter(img_data, ww, wc)
        else:
            img_data = convert_defalut_ww_wc(img_data)
    else:
        img_data = convert_defalut_ww_wc(img_data)

    if dcm_info['photo_metric_interpretation'].upper().find('MONOCHROME1') >= 0:
        img_data = (255 - img_data)

    return img_data


def get_pixeldata(dicom_dataset):
    """If NumPy is available, return an ndarray of the Pixel Data.
    Raises
    ------
    TypeError
        If there is no Pixel Data or not a supported data type.
    ImportError
        If NumPy isn't found
    NotImplementedError
        if the transfer syntax is not supported
    AttributeError
        if the decoded amount of data does not match the expected amount
    Returns
    -------
    numpy.ndarray
       The contents of the Pixel Data element (7FE0,0010) as an ndarray.
    """
    if (dicom_dataset.file_meta.TransferSyntaxUID not in
            NumpySupportedTransferSyntaxes):
        raise NotImplementedError("Pixel Data is compressed in a "
                                  "format pydicom does not yet handle. "
                                  "Cannot return array. Pydicom might "
                                  "be able to convert the pixel data "
                                  "using GDCM if it is installed.")

    #dicom_dataset.
    if 'PixelData' not in dicom_dataset:
        raise TypeError("No pixel data found in this dataset.")

    # Make NumPy format code, e.g. "uint16", "int32" etc
    # from two pieces of info:
    # dicom_dataset.PixelRepresentation -- 0 for unsigned, 1 for signed;
    # dicom_dataset.BitsAllocated -- 8, 16, or 32
    if dicom_dataset.BitsAllocated == 1:
        # single bits are used for representation of binary data
        format_str = 'uint8'
    elif dicom_dataset.PixelRepresentation == 0:
        format_str = 'uint{}'.format(dicom_dataset.BitsAllocated)
    elif dicom_dataset.PixelRepresentation == 1:
        format_str = 'int{}'.format(dicom_dataset.BitsAllocated)
    else:
        format_str = 'bad_pixel_representation'
    try:
        numpy_dtype = np.dtype(format_str)
    except TypeError:
        msg = ("Data type not understood by NumPy: "
               "format='{}', PixelRepresentation={}, "
               "BitsAllocated={}".format(
                   format_str,
                   dicom_dataset.PixelRepresentation,
                   dicom_dataset.BitsAllocated))
        raise TypeError(msg)

    if dicom_dataset.is_little_endian != sys_is_little_endian:
        numpy_dtype = numpy_dtype.newbyteorder('S')

    pixel_bytearray = dicom_dataset.PixelData

    if dicom_dataset.BitsAllocated == 1:
        # if single bits are used for binary representation, a uint8 array
        # has to be converted to a binary-valued array (that is 8 times bigger)
        try:
            pixel_array = np.unpackbits(
                np.frombuffer(pixel_bytearray, dtype='uint8'))
        except NotImplementedError:
            # PyPy2 does not implement numpy.unpackbits
            raise NotImplementedError(
                'Cannot handle BitsAllocated == 1 on this platform')
    else:
        pixel_array = np.frombuffer(pixel_bytearray, dtype=numpy_dtype)
    length_of_pixel_array = pixel_array.nbytes
    expected_length = dicom_dataset.Rows * dicom_dataset.Columns
    if ('NumberOfFrames' in dicom_dataset and
            dicom_dataset.NumberOfFrames > 1):
        expected_length *= dicom_dataset.NumberOfFrames
    if ('SamplesPerPixel' in dicom_dataset and
            dicom_dataset.SamplesPerPixel > 1):
        expected_length *= dicom_dataset.SamplesPerPixel
    if dicom_dataset.BitsAllocated > 8:
        expected_length *= (dicom_dataset.BitsAllocated // 8)
    padded_length = expected_length
    if expected_length & 1:
        padded_length += 1
    if length_of_pixel_array != padded_length:
        raise AttributeError(
            "Amount of pixel data %d does not "
            "match the expected data %d" %
            (length_of_pixel_array, padded_length))
    if expected_length != padded_length:
        pixel_array = pixel_array[:expected_length]
    if should_change_PhotometricInterpretation_to_RGB(dicom_dataset):
        dicom_dataset.PhotometricInterpretation = "RGB"
    if dicom_dataset.Modality.lower().find('ct') >= 0:  # CT图像需要得到其CT值图像
        pixel_array = pixel_array * dicom_dataset.RescaleSlope + dicom_dataset.RescaleIntercept  # 获得图像的CT值
    pixel_array = pixel_array.reshape(dicom_dataset.Rows, dicom_dataset.Columns*dicom_dataset.SamplesPerPixel)
    return pixel_array, dicom_dataset.Rows, dicom_dataset.Columns


def convert_defalut_ww_wc(img_data):
    img_temp = img_data
    img_temp = (np.maximum(img_temp, 0) / img_temp.max()) * 255.0
    return np.uint8(img_temp)


def setDicomWinWidthWinCenter(img_data, ww, wc):
    img_temp = img_data
    img_temp.flags.writeable = True
    min_v = (2 * wc - ww) / 2.0 + 0.5
    max_v = (2 * wc + ww) / 2.0 + 0.5
    factor = 255.0 / (max_v - min_v)

    img_temp = (img_temp - min_v) * factor

    min_index = img_temp < 0
    img_temp[min_index] = 0
    max_index = img_temp > 255
    img_temp[max_index] = 255
    return np.uint8(img_temp)

def getPixelsHu(slices):
    """
    convert 2d dicom slices into CT value(int16)
    :param slices:2d dicom slices
    :return:3D CT image(n*w*h), spacing
    """
    image = np.stack([s.pixel_array for s in slices])



    # Convert to int16 (from sometimes int16),
    # should be possible as values should always be low enough (<32k)
    image = image.astype(np.int16)

    # Convert to Hounsfield units (HU)
    for slice_number in range(len(slices)):
        intercept = slices[slice_number].RescaleIntercept
        slope = slices[slice_number].RescaleSlope

        if slope != 1:
            image[slice_number] = slope * image[slice_number].astype(np.float64)
            image[slice_number] = image[slice_number].astype(np.int16)

        image[slice_number] += np.int16(intercept)

    return np.array(image, dtype=np.int16), np.array([slices[0].SliceThickness] + slices[0].PixelSpacing._list,
                                                     dtype=np.float32)

def printScanSomeInfo(path):
    slices = []
    namedict = {}
    if isinstance(path, str):
        dcms = glob.glob(path + "/*.dcm")
        for s in dcms:
            ds = pydicom.dcmread(s)
            ds.decompress()
            print("StudyID: %s, SIUID: %s, SerireInsUID: %s"%(ds.StudyID, ds.StudyInstanceUID, ds.SeriesInstanceUID))

def main():
    import matplotlib.pyplot as plt
    path = '/home/blake/medical/DR/QC/gdDR/data/77/63186359.dcm'
    path = '../data/tiwaiyiwu.dcm'
    ds_info = read_CR_dcm(path)
    plt.imshow(ds_info['image'], cmap=plt.cm.gray)
    plt.show()


if __name__ == '__main__':
    main()



